#!/usr/bin/env python3
# coding=utf8

import requests
from datetime import datetime, timedelta
import json
import time
import sys
import argparse
import subprocess
import logging
import xml.etree.ElementTree as ET

# Domain = "http://openapi.zhiyan.sg.woa.com" # 海外
Domain = "http://openapi.zhiyan.woa.com" # 国内

Token = "2b46dcc220c371a3e1e0308489668265" # 认证token
Staffname = "rexjliu"
projectname = "osg" # 项目在智妍的项目名
project_ch_name = "OSG" # 项目在智妍的项目名
appmark = "1168_17986_OSGame"
AppName = "OSGame"

# 接入名称映射，日志汇-> 接入管理-> 接入名称
LOG_ACCESS_NAME = {
    "prod": "osgame_prod",
    "ce2": "osgame_ce",
    "pre": "osgame_pre",
    "opcheck": "osgame_prod_opcheck",
}

# 日志汇节点日志数据检查排除列表IP
EXCLUDE_LOG_IPS = ["9.150.19.50", "9.150.12.212", "30.49.245.203", "9.150.123.146", "9.150.111.5"]

MaxRetryCount = 10

logger = logging.getLogger()

Headers = {
    "token": Token,
    "staffname": Staffname,
    "projectname": projectname,
    "Content-Type": "application/json"
}

def req():
    requests.post()


class XmlParser:
    def __init__(self, filepath, fileType):
        self.filepath = filepath
        self.mytree = ET.parse(self.filepath)
        self.myroot = self.mytree.getroot()
        self.filetype = fileType
        self.myroot = self.check_file()

    def check_file(self):
        success = False
        CheckTag = ""
        if self.filetype == "host":
            CheckTag = "HostTab"
        elif self.filetype == "proc":
            CheckTag = "cluster"
        elif self.filetype == "proc_deploy":
            CheckTag = "ClusterDeploy"
        elif self.filetype == "bus_relation":
            CheckTag = "TbusConf"
        else:
            raise Exception("unsuport type: %s"%(self.filetype))

        for i in self.myroot:
            if i.tag == CheckTag:
                success = True
                return i

        if not success:
            raise Exception("not found type: %s tag: %s"%(self.filetype, CheckTag))

# 获取infosvr的请求次数
def req_infosvr(env, ips):
    sec_lvl_name = "bingoTransport"
    metric_name = "bingo.http.TransportHandle_cnt"
    tag_set = [
        {"key": "funcid", "value": ["104"]} # 106是dirsvr，104是infosvr
    ]
    return req_query_data(env, ips, sec_lvl_name, metric_name, tag_set)

# 获取dirsvr请求次数
def req_dirsvr(env, ips):
    sec_lvl_name = "dirsvr"
    metric_name = "dirsvr.getloginvport.successcnt"
    tag_set = [
    ]
    return req_query_data(env, ips, sec_lvl_name, metric_name, tag_set)


def req_query_data(env, ips, sec_lvl_name, metric_name, tag_set, retrycount=0):
    url_path = "/monitor/v2/api/chart/info/query"
    now = datetime.now()
    data = {
        "app_mark": appmark,
        "sec_lvl_name": sec_lvl_name,
        "metric_name": metric_name,
        "env": env,
        "is_english": "yes",
        "is_together": False,
        "instance_mark": ips,
        "begin_time": (now - timedelta(minutes=2)).strftime("%Y-%m-%d %H:%M:%S"),
        "end_time": now.strftime("%Y-%m-%d %H:%M:%S"),
        "tag_set": tag_set
    }

    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
        logger.error("Error: get failed with http status code: %d"%(rsp.status_code))
        if retrycount > MaxRetryCount:
            raise Exception("query failed")
        time.sleep(1)
        return req_query_data(env, ips, retrycount+1)

    rsp_data = rsp.json()
    if rsp_data["code"] != 0:
        raise Exception("query failed: %s"%(rsp_data["msg"]))
    logger.info(json.dumps(rsp_data))
    for chart in rsp_data["data"]["chart_info"]:
        # logger.info(chart["title"], chart["detail_data_list"])
        return chart["detail_data_list"]

def add_ip(env, ips):
    url_path = "/infra/v1/infra-inst/upsert"
    data = {
            "instances": [
            ]
    }
    for ip in ips:
        data["instances"].append({
            "app_id": get_appid(),
            "device_ip": ip,
            "env_name": env,
            "identifier": ip,
            "env_id": get_env_info(env),
            "host_type": 1,
        })
    
    logger.info("data: %s", data)
    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
        raise Exception("Error: get failed with http status code: %d"%(rsp.status_code))
    
    rsp_data = rsp.json()
    logger.info("get resp: %s", rsp_data)
    if rsp_data["code"] != 0:
        raise Exception("add ip code not 0, failed: %s"%(rsp_data))

# 删除监控宝IP
def del_ip(env, ips):
    url_path = "/infra/v1/infra-inst/delete"
    data = {
            "instances": [
            ]
    }
    appid = get_appid()
    envid = get_env_info(env)
    for ip in ips:
        data["instances"].append({
            "app_id": appid,
            "env_id": envid,
            "identifier": ip,
        })
    logger.info("del_ip data: %s", data)
    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
        raise Exception("del_ip status code: %d, content:%s"%(rsp.status_code, rsp.content))
    
    rsp_data = rsp.json()
    logger.info("get resp: %s", rsp_data)

    if rsp_data["code"] != 0:
        raise Exception("del ip code not 0, failed: %s"%(rsp_data))

# 获取实例列表
def query_ips(env, ips):
    url_path = "/infra/v1/infra-inst/query"
    appid = get_appid()

    page_num = 1
    page_size = 50
    total_count = 0
    first_stage = True
    results = []
    
    instance_filter = [
            {
                "key": "basic.app_id",
                "op": "eq",
                "value": appid,
            },
            {
                "key": "basic.env_name",
                "op": "eq",
                "value": env,
            },
        ]
    
    if ips:
        instance_filter.append({
                "key": "device.ip",
                "op": "in",
                "value": ips,
        })

    logger.info("req data %s"%(instance_filter))
    while True:
        if not first_stage and len(results) >= total_count:
            break
        first_stage = False
        data = {
            "page": page_num,
            "per_page": page_size,
            "instance_filter": json.dumps(instance_filter)
        }

        rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
        if rsp.status_code != 200:
            raise Exception("del_ip status code: %d, content:%s"%(rsp.status_code, rsp.content))
        
        rsp_data = rsp.json()
        logger.info("get query ip info: %s"%(rsp_data))
        if rsp_data["code"] != 0:
            raise Exception("code not 0")

        results += rsp_data["data"]["items"]
        total_count = rsp_data["data"]["total"]
        page_num += 1

    return results

# env_type	环境类型. 可选：prod/test/pre_release/dev
def get_env_info(env):
    url_path = "/infra/v1/infra-env/query"
    data = {
        "en_name": env,
        "page": 1,
        "per_page": 1000,
        "project_id": get_project(),
    }

    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
        raise Exception("Error: get failed with http status code: %d"%(rsp.status_code))
    
    rsp_data = rsp.json()
    logger.info("get env resp: %s", json.dumps( rsp_data))
    return rsp_data["data"]["items"][0]["id"]

def get_project():
    url_path = "/infra/v1/infra-project/list"
    data = {
        "name": project_ch_name,
        "page": 1,
        "per_page": 1000,
    }
    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
        raise Exception("Error: get failed with http status code: %d"%(rsp.status_code))
    
    rsp_data = rsp.json()
    logger.info("get project resp: %s", json.dumps( rsp_data))
    return rsp_data["data"][0]["id"]

# 获取日志汇接入列表
def get_log_access_list(env):
    url_path = "/log/v3/dataflow/list"
    data = {
        "projectId": get_project(),
        "applicationId": -1,
        "name": LOG_ACCESS_NAME[env],
        "environment": env,
    }
    logger.info("req data: %s", data)
    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
        raise Exception("get log accesslist status code not 200: %d, content: %s"%(rsp.status_code, rsp.content))
    
    rsp_data = rsp.json()
    if rsp_data["code"] != 0:
        raise Exception("get log access list code not 0: %s"%(rsp_data))
    
    logger.info("get access list: %s", json.dumps(rsp_data))
    for i in rsp_data["data"]["list"]:
        if i["name"] == LOG_ACCESS_NAME[env]:
            return i["id"]
    raise Exception("get access log id failed")

# 获取日志汇接入IP列表
def get_log_instances_status(env):
    url_path = "/log/v3/dataflow/collect/delivery/list"
    access_id = get_log_access_list(env)
    projectid = get_project()
    totalCount = 0
    first = True
    result_list = []
    page = 1
    pageSize = 100
    while True:
        if not first and len(result_list) >= totalCount:
            break
        first = False
        data = {
            "pageNum": page,
            "pageSize": pageSize,
            # "version": 11,
            "projectId": projectid,
            # "dataflowName": LOG_ACCESS_NAME[env],
            "dataflowId": access_id,
        }
        rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
        if rsp.status_code != 200:
            raise Exception("req log list failed, status code not 200: %d"%(rsp.status_code))

        rsp_data = rsp.json()
        if rsp_data["code"] != 0:
            raise Exception("req log list faled: code: %d, msg: %s"%(rsp_data["code"], rsp_data["msg"]))
        page += 1
        totalCount = rsp_data["data"]["totalCount"]
        result_list += rsp_data["data"]["list"]
        logger.info("get data: %s", rsp_data)
    return result_list

def check_rizhihui_agent_logcount_realtime(env, ignoreerror):
    url_path = "/log/v2/search/group"
    s = datetime.strptime(datetime.now().strftime("%Y%m%d"), "%Y%m%d")
    startts = time.mktime(s.timetuple())
    data = {
            "startTime": (startts - 60 * 60 * 0.6)*1000 ,
            "endTime": startts * 1000,
            "dataflowName": LOG_ACCESS_NAME[env],
            "groupField": "@host",
            "storageType": "elasticsearch",
            "limit": 10000
        }
    
    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
            raise Exception("req log list failed, status code not 200: %d, rsp: %s"%(rsp.status_code, rsp.text))

    rsp_data = rsp.json()
    if rsp_data["code"] != 0:
        raise Exception("req log list faled: code: %d, msg: %s"%(rsp_data["code"], rsp_data["msg"]))
    print("get report data:" , json.dumps( rsp_data, indent=4))

    host_xml = XmlParser("/data/home/user00/apps/tcm/cfg/host.xml", "host")

    fail_hosts = []
    for i in host_xml.myroot:
        ip = i.attrib.get("InnerIP")
        check = False
        for j in rsp_data['data']:
            if not check:
                for ip1, count in j.items():
                    if ip == ip1 and count > 0:
                        check = True
                        break
            else:
                break
        if not check:
            fail_hosts.append(ip)
    print("<SOPS_VAR>failip:%s</SOPS_VAR>"%(",".join(fail_hosts)))
    if len(fail_hosts) > 0 and not ignoreerror:
        raise Exception("ips %s not get log"%(",".join(fail_hosts)))

#获取日志汇日志上报数量
def check_rizhihui_agent_logcount(env, ignoreerror):
    url_path = "/log/v3/agent/statistics"

    access_id = get_log_access_list(env)

    s = datetime.strptime(datetime.now().strftime("%Y%m%d"), "%Y%m%d")
    startts = time.mktime(s.timetuple())

    data = {
        "projectId": get_project(),
        "dataflowIdList": [access_id],
        "startTime": startts * 1000,
        "endTime": (startts + 60 * 60 * 24)*1000,
        "dataType": "REPORT"
    }

    print("req data: ", data)
    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
            raise Exception("req log list failed, status code not 200: %d"%(rsp.status_code))

    rsp_data = rsp.json()
    if rsp_data["code"] != 0:
        raise Exception("req log list faled: code: %d, msg: %s"%(rsp_data["code"], rsp_data["msg"]))
    print("get report data:" , json.dumps( rsp_data, indent=4))
    error_ips = []
    for i in rsp_data["data"]["data"]:
        if int(i["count"]) == 0 and i["host"] not in EXCLUDE_LOG_IPS:
            print("ERROR: host %s, log count is 0"%(i["host"]))
            error_ips.append(i["host"])

    print("<SOPS_VAR>failip:%s</SOPS_VAR>"%(",".join(error_ips)))
    if len(error_ips) > 0 and not ignoreerror:
            raise Exception("ips %s not get log"%(",".join(error_ips)))

def update_log_instances(env, add_ips):
    url_path = "/log/v3/dataflow/collect/delivery/register"
    data ={
        "id": get_log_access_list(env),
        "instances": add_ips,
    }

    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
        raise Exception("register log ips req failed: status code: %d, msg: %s"%(rsp.status_code, rsp.content))

    rsp_data = rsp.json()
    logger.info("update log instances result: %s"%(rsp_data))
    if rsp_data["code"] != 0:
        raise Exception("update log instances failed: %s"%(rsp_data["msg"]))

    success_instances = rsp_data["data"]["instances"]
    failed = []
    for ip in add_ips:
        if ip not in success_instances:
            failed.append(ip)

    if len(failed) > 0:
        raise Exception("had ip not in the result: %s"%(failed))

def delete_log_instances(env, ipsstr):
    url_path = "/log/v3/dataflow/collect/delivery/delete"
    data = {
        "dataflowId": get_log_access_list(env),
        "ip": ipsstr,
    }

    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
        raise Exception("delete_log_instances req failed: status code: %d, msg: %s"%(rsp.status_code, rsp.content))

    rsp_data = rsp.json()
    logger.info("delete_log_instances result: %s"%(rsp_data))
    if rsp_data["code"] != 0:
        raise Exception("delete_log_instances failed: %s"%(rsp_data["msg"]))

    if not rsp_data["data"]:
        raise Exception("delete_log_instances failed")

def get_appid():
    url_path = "/infra/v1/infra-app/list"
    data  = {
        "name": AppName,
        "page": 1,
        "per_page": 1000,
        "project_id": get_project(),
    }
    rsp = requests.post(Domain + url_path, data=json.dumps(data), headers=Headers)
    if rsp.status_code != 200:
        logger.error("Error: get failed with http status code: %d"%(rsp.status_code))
    
    rsp_data = rsp.json()
    logger.info("get appinfo resp: %s", json.dumps( rsp_data))
    return rsp_data["data"][0]["id"]

def cmd(cmdstr):
    c = subprocess.run(cmdstr, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out = c.stdout.decode("utf-8")
    err = c.stderr.decode("utf-8")
    if c.returncode != 0:
        raise Exception("execute cmd: %s , failed: %s"%(cmdstr, err))
    logger.info("execute cmd: %s, stdout: %s, stderr: %s"%(cmdstr, out, err))
    return out

def checkzhiyan(env, ips):
    data = query_ips(env, ips)

    for i in data:
        ip = i["device"]["ip"]
        if ip in ips:
            ips.remove(ip)

    if len(ips) > 0:
        raise Exception("not get the ips: %s"%(",".join(ips)))

def checklog(env, ips, ignoreerror):
    results = get_log_instances_status(env)
    errors = []
    for i in results:
        if i["ip"] not in ips:
            continue
        print('zhiyan_log_status{log_agent_ip="%s"} %d'%(i["ip"], i["status"]))
        if i["status"] != 2:
            errors.append(i)
        ips.remove(i["ip"])

    for i in ips:
        print('zhiyan_log_status{log_agent_ip="%s"} -1'%(i))
    if len(errors) > 0 and not ignoreerror:
        raise Exception("some ips status not online %s"%(errors))
    if len(ips) > 0 and not ignoreerror:
        raise Exception("some ips not in the log instances list: %s"%(ips))

def parse():
    parser = argparse.ArgumentParser("zhiyan tools")
    parser.add_argument("-a", "--action", required=True, choices=["addzhiyan", "delzhiyan", "addlog", "dellog", "checkzhiyan", "checklog", "logreport"], help="need the action")
    parser.add_argument("-e", "--env", required=True, help="must input the env")
    parser.add_argument("-i", "--ips", default="", help="the ips, multipule ip use , connect")
    parser.add_argument("-g", "--ignoreerror", default=False, action="store_true", help="only return status 0")
    return parser.parse_args()

if __name__ == "__main__":
    args = parse()

    level = logging.INFO
    if args.ignoreerror:
        level = logging.ERROR
    logger.setLevel(level)
    handler = logging.StreamHandler(sys.stdout)
    handler.setLevel(level)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    
    if args.action == "addzhiyan":  # 添加监控宝IP
        if args.ips.strip() == "":
            raise Exception("must give add ips")
        ips = args.ips.strip().split(",")
        add_ip(args.env, ips)

        checkzhiyan(args.env, ips)
    elif args.action == "checklog":  # 检查日志汇的agent状态
        if args.ips == "":
            ips = cmd("""grep InnerIP /data/home/user00/apps/tcm/cfg/host.xml | awk -F 'InnerIP=\"' '{print $2}' | cut -d '"' -f 1""").split()
        else:
            ips = args.ips.split(",")

        checklog(args.env, ips, args.ignoreerror)
        
    elif args.action == "addlog":
        if args.ips.strip() == "":
            raise Exception("must give add ips")
        ips = args.ips.strip().split(",")
        update_log_instances(args.env, ips)
        logger.info("wait 120 second to check agent online")
        time.sleep(120)
        checklog(args.env, ips, False)

    elif args.action == "dellog":
        if args.ips.strip() == "":
            raise Exception("must give add ips")
        delete_log_instances(args.env, args.ips)

    elif args.action == "delzhiyan":
        if args.ips.strip() == "":
            raise Exception("must give add ips")
        del_ip(args.env, args.ips.strip().split(","))

    elif args.action == "checkzhiyan":
        if args.ips == "":
            ips = cmd("""grep InnerIP /data/home/user00/apps/tcm/cfg/host.xml | awk -F 'InnerIP=\"' '{print $2}' | cut -d '"' -f 1""").split()
        else:
            ips = args.ips.split(",")

        checkzhiyan(args.env, ips)
    elif args.action == "logreport":
        check_rizhihui_agent_logcount_realtime(args.env, args.ignoreerror)
    else:
        raise Exception("not surport this action")